"""
Post-processing script for HEMC validation.
处理 Geant4 scoring mesh, 要求必须是 1 Dimension 的,
从 .npz 文件读取 quantity 并作图.
"""

import glob
import re
from pathlib import Path

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np

from ..config.models import HEMCConfig
from ..config.router import router

_MODULENAME = Path(__file__).stem
_DUMP_PATTERN = '_*.npz'


def parse_npz(path: Path) -> dict:
    """
    Parse Geant4 mesh .npz file.

    Args:
        path: Path to .npz file

    Returns:
        # iZ, iPHI, iR, total(value) [MeV], total(val^2), entry
        Dictionary mapping iZ to (sum_val, sum_val2, n)
    """
    # with np.load(path, mmap_mode='r', allow_pickle=False) as mynpz:
    with np.load(path, mmap_mode='r') as mynpz:
        iZ, value, val2, entry = (
            mynpz["idx0"],
            mynpz["value"],
            mynpz["val2"],
            mynpz["entry"],
        )
        # 打包成 150 行，每行 4 个数
        data = np.column_stack((iZ, value, val2, entry))
        return data


def calc_mean_var(dat: dict) -> tuple:
    """
    Calculate mean and variance from parsed .npz data.

    Args:
        dat: Dictionary from parse_npz

    Returns:
        Tuple of (iz_arr, mean_arr, var_arr, sigma_arr)
    """
    # izs = sorted(dat.keys())
    izs = dat[:, 0].astype(np.int32)
    v_means, v_vars = [], []
    for row in dat:
        iz, sv, sv2, n = row
        v_m = 0.0
        v_var = 0.0
        if n != 0:
            v_m = sv / n
            v2_m = sv2 / n
            if n != 1:
                v_var = (v2_m - v_m * v_m) / (n - 1)  # Population variance
        v_means.append(v_m)
        v_vars.append(v_var)
    return np.array(izs), np.array(v_means), np.array(v_vars), np.sqrt(v_vars)


def extract_prefix_energy(filename: str) -> tuple:
    """
    Extract prefix and energy from filename.

    Args:
        filename: Name of the CSV file

    Returns:
        Tuple of (prefix, energy) or (None, None) if not found
    """
    pattern = r"(\w+)_(\d+\.?\d*)MeV"
    match = re.match(pattern, filename)

    if match:
        prefix = match.group(1)  # e.g., IDD, LDose, LETd, eDep
        energy = match.group(2)  # e.g., 110.0, 104.0
        return prefix, energy
    else:
        return None, None


def plot_summary(prefix: str, file_list: list, config: HEMCConfig, save: bool = True) -> None:
    """
    Create a summary plot for a given prefix with all energy levels.

    Args:
        prefix: The prefix (e.g., IDD, LDose, LETd, eDep)
        file_list: List of CSV files with the same prefix
        config: HEMC configuration
        save: Whether to save the plot
    """
    plt.figure(figsize=config.plot.figsize)

    # Define a color map for different energies
    # colors = plt.cm.get_cmap('viridis')(np.linspace(0, 1, len(file_list)))
    lenfl = len(file_list)
    lsp = np.linspace(0, 1, lenfl)
    colors = mpl.colormaps['viridis'](lsp)

    for i, f in enumerate(file_list):
        csv_file = Path(f)
        basename = csv_file.stem

        # Extract energy from filename
        _, energy = extract_prefix_energy(basename)

        dat = parse_npz(csv_file)
        iz, mean, var, sigma = calc_mean_var(dat)

        # Plot mean values with different colors
        plt.plot(iz, mean, label=f"{prefix} @ {energy} MeV", color=colors[i], lw=1.2)

        # Optionally also plot the error bands
        plt.fill_between(iz, mean - sigma, mean + sigma, color=colors[i], alpha=0.1)

    # Set labels and title
    ylabel = config.plot.ylabel_dict.get(prefix, prefix)
    plt.xlabel(config.plot.xlabel)
    plt.ylabel(ylabel)
    plt.title(f"{prefix} Summary (All Energies)")
    plt.tight_layout()
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')  # Place legend outside plot to avoid overlap

    if save:
        assert config.project_paths.resultdir is not None
        outpng = config.project_paths.resultdir / f"{prefix}_summary.png"
        plt.savefig(outpng, dpi=config.plot.save_dpi, bbox_inches='tight')
        print(f"Saved summary plot -> {outpng}")
    plt.close()


def plot_one(npz_file: Path, config: HEMCConfig, save: bool = True) -> None:
    """
    Process single npz file and generate plot.

    Args:
        npz_file: Path to input npz file
        config: HEMC configuration
        save: Whether to save the plot
    """
    basename = npz_file.stem

    # Extract scorer name and energy from filename using regex
    # Pattern matches: prefix_energyMeV (e.g., IDD_110.00MeV, LDose_104.00MeV)
    pattern = r"(\w+)_(\d+\.?\d*)MeV"
    match = re.match(pattern, basename)

    if match:
        scorer = match.group(1)  # e.g., IDD, LDose, LETd, eDep
        energy = match.group(2)  # e.g., 110.00, 104.00
    else:
        scorer = "quantity"
        energy = "unknown"

    ylabel = config.plot.ylabel_dict.get(scorer, scorer)

    dat = parse_npz(npz_file)
    iz, mean, var, sigma = calc_mean_var(dat)

    plt.figure(figsize=config.plot.figsize)
    plt.fill_between(iz, mean - sigma, mean + sigma, color="C0", alpha=0.25, label="±1σ")
    plt.plot(iz, mean, lw=1.2, color="C0", label="Mean")
    plt.xlabel(config.plot.xlabel)
    plt.ylabel(ylabel)
    plt.title(f"{scorer} @ {energy} MeV")
    plt.tight_layout()
    plt.legend()

    if save:
        assert config.project_paths.resultdir is not None
        outpng = config.project_paths.resultdir / f"{basename}.png"
        plt.savefig(outpng, dpi=config.plot.save_dpi)
        print(f"Saved  -> {outpng}")
    plt.close()


@router.register()
def plotnpz1D(config: HEMCConfig):
    """Main entry point for post-processing."""

    # Use whitelist for specific prefixes instead of processing all CSV files
    # prefixes = ["eDep", "IDD", "LDose", "LETd"]
    prefixes = ["eDep", "IDD", "LETd"]
    all_files = []

    assert config.project_paths.g4outdir is not None
    for prefix in prefixes:
        pattern = str(config.project_paths.g4outdir / f"{prefix}{_DUMP_PATTERN}")
        files = glob.glob(pattern)
        all_files.extend(files)

    if not all_files:
        print("No CSV files found matching the prefixes: eDep, IDD, LDose, LETd!")
        return

    # Group files by prefix for summary plots
    prefix_groups = {p: [] for p in prefixes}
    for f in all_files:
        filename = Path(f).stem
        # Extract prefix directly from the beginning of the filename
        for prefix in prefixes:
            if filename.startswith(prefix + "_"):
                prefix_groups[prefix].append(f)
                break

    # Create individual plots for each file
    print("Creating individual plots...")
    for f in all_files:
        print(f"[plotnpz] Processing {f} ...")
        plot_one(Path(f), config, save=True)

    # Create summary plots for each prefix with all energy levels
    if False:
        print("Creating summary plots...")
        for prefix in prefixes:
            file_list = prefix_groups[prefix]
            if file_list:  # Only create summary plot if there are files for this prefix
                print(f"Creating summary plot for {prefix} with {len(file_list)} energy levels...")
                plot_summary(prefix, file_list, config, save=True)


if __name__ == "__main__":
    pass
